# -*- coding: utf-8 -*-

'''
Created on 2018年9月1日

@author: Dergen Lee

'''

import os
from tensorflow import keras
import gensim
import gzip
import numpy as np
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

gnews = 'GoogleNews-vectors-negative300.bin'

unzipped = os.path.join(os.environ['GENSIM_DATASETS'], 'GoogleNews', gnews)

print('Getting pretrained model: GoogleNews-vectors-negative300.bin')
if not os.path.isfile(unzipped):
    path = keras.utils.get_file(os.path.join(os.environ['GENSIM_DATASETS'], 'GoogleNews', gnews + '.gz'),
              'https://s3.amazonaws.com/dl4j-distribution/%s.gz' % gnews)
    with open(unzipped, 'wb') as fout:
        g = gzip.GzipFile(mode="rb", fileobj=open(path, 'rb'))
        fout.write(g.read())

print('Loading a pretrained word embedding model')
model = gensim.models.KeyedVectors.load_word2vec_format(unzipped, binary=True)

print('Get most similar of china')
likes = model.most_similar(positive=['china'])
print(likes)


def A_is_to_B_as_C_is_to(a, b, c, topn=1):
    a, b, c = map(lambda x:x if type(x) == list else [x], (a, b, c))
    
    # find res where v[b]-v[a]+v[c]-v[res] is minumum
    res = model.most_similar(positive=b + c, negative=a, topn=topn)
    
    if len(res):
        if topn == 1:
            return res[0][0]
        return [x[0] for x in res]
    return None


print(A_is_to_B_as_C_is_to("man", "woman", "boy", topn=1))

for country in 'Italy', 'France', 'India', 'China':
    print('%s is the capital of %s' % 
          (A_is_to_B_as_C_is_to('Germany', 'Berlin', country), country))
    
for company in 'Google', 'IBM', 'Boeing', 'Microsoft', 'Samsung':
    products = A_is_to_B_as_C_is_to(
      ['Starbucks', 'Apple'], ['Starbucks_coffee', 'iPhone'], company, topn=3)
    print('%s -> %s' % (company, ', '.join(products)))    

beverages = ['espresso', 'beer', 'vodka', 'wine', 'cola', 'tea']
countries = ['Italy', 'Germany', 'Russia', 'France', 'USA', 'India']
sports = ['soccer', 'handball', 'hockey', 'cycling', 'basketball', 'cricket']

items = beverages + countries + sports

item_vectors = [(item, model[item])
                    for item in items
                    if item in model]

vectors = np.asarray([x[1] for x in item_vectors])
lengths = np.linalg.norm(vectors, axis=1)
norm_vectors = (vectors.T / lengths).T
tsne = TSNE(n_components=2,
            perplexity=10,
            verbose=2).fit_transform(norm_vectors)
            
x = tsne[:, 0]
y = tsne[:, 1]

fig, ax = plt.subplots()
ax.scatter(x, y)

for item, x1, y1 in zip(item_vectors, x, y):
    ax.annotate(item[0], (x1, y1))

plt.show()

